import argparse
import sys
import time
from dataclasses import dataclass

import cv2 as cv
import numpy as np

# ROS imports (expected on Jetson with ROS installed)
try:
    import rospy
    from geometry_msgs.msg import Twist
except Exception as ros_import_error:  # noqa: F841
    rospy = None
    Twist = None

# Optional UDP control for mode switching if available
try:
    from DogMotion import ControlRobot  # type: ignore
except Exception:
    ControlRobot = None  # type: ignore


@dataclass
class NavParams:
    target_id: int = 1
    marker_size_m: float = 0.10  # 10 cm (for reference; we use pixel-size stop)
    left_right_thr_px: int = 10
    forward_center_thr_px: int = 30
    stop_side_px: int = 160  # stop when one side length >= this pixel size
    linear_speed: float = 0.25
    angular_speed: float = 0.6
    low_linear_speed: float = 0.10
    low_angular_speed: float = 0.30
    cmd_rate_hz: float = 20.0


def build_aruco_detector() -> tuple:
    """Create an ArUco detector supporting different OpenCV versions.

    Returns a tuple (mode, dictionary, detector_or_params)
    - mode = 'new' uses ArucoDetector
    - mode = 'legacy' uses detectMarkers with parameters
    """
    try:
        dictionary = cv.aruco.getPredefinedDictionary(cv.aruco.DICT_4X4_50)
        detector = cv.aruco.ArucoDetector(dictionary)
        return "new", dictionary, detector
    except AttributeError:
        dictionary = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_50)
        params = cv.aruco.DetectorParameters_create()
        return "legacy", dictionary, params


def compute_marker_side_length_px(corner_array: np.ndarray) -> float:
    """Compute average side length in pixels for a single marker corners array.

    corner_array shape is (4, 1, 2) or (4, 2). We will squeeze to (4, 2).
    """
    corners = np.squeeze(corner_array)  # -> (4, 2)
    p0, p1, p2, p3 = corners
    d01 = np.linalg.norm(p1 - p0)
    d12 = np.linalg.norm(p2 - p1)
    d23 = np.linalg.norm(p3 - p2)
    d30 = np.linalg.norm(p0 - p3)
    return float((d01 + d12 + d23 + d30) / 4.0)


def get_marker_center_x(corner_array: np.ndarray) -> float:
    corners = np.squeeze(corner_array)
    return float(np.mean(corners[:, 0]))


def publish_twist(pub, linear_x: float, angular_z: float):
    msg = Twist()
    msg.linear.x = float(linear_x)
    msg.angular.z = float(angular_z)
    pub.publish(msg)


def set_autonomous_low_pose_if_possible():
    if ControlRobot is None:
        return
    try:
        controller = ControlRobot()
        # Switch to autonomous mode
        controller.autonomous_mode()
        # Enter move mode with low-speed gait and lower the body height to approximate a prone/crawl posture
        controller.move_mode()
        controller.flat_low_speed_gait()
        controller.adjust_height(-20000)
    except Exception as e:
        print(f"[WARN] Failed to set autonomous low pose via UDP: {e}")


def switch_back_to_manual_if_possible():
    if ControlRobot is None:
        return
    try:
        controller = ControlRobot()
        controller.manual_mode()
    except Exception as e:
        print(f"[WARN] Failed to switch back to manual via UDP: {e}")


def main():
    parser = argparse.ArgumentParser(description="Jetson NX ArUco navigation publishing to /cmd_vel")
    parser.add_argument("--camera", default="/dev/video4", help="Video device path, e.g., /dev/video3")
    parser.add_argument("--gui", action="store_true", help="Enable GUI display")
    parser.add_argument("--low", action="store_true", help="Low power mode (slower speeds)")
    parser.add_argument("--target-id", type=int, default=1, help="ArUco target ID")
    parser.add_argument("--stop-px", type=int, default=160, help="Stop when marker side length reaches this pixel size")
    parser.add_argument("--align-thr", type=int, default=30, help="Forward alignment pixel threshold")
    parser.add_argument("--lr-thr", type=int, default=10, help="Left-right pixel threshold")
    args = parser.parse_args()

    if rospy is None or Twist is None:
        print("[ERROR] ROS is not available. Please ensure rospy and geometry_msgs are installed in this environment.")
        sys.exit(1)

    # ROS setup
    rospy.init_node("jetson_aruco_nav", anonymous=True)
    cmd_pub = rospy.Publisher("/cmd_vel", Twist, queue_size=10)
    rate = rospy.Rate(20)

    # Navigation parameters
    params = NavParams(
        target_id=args.target_id,
        stop_side_px=args.stop_px,
        forward_center_thr_px=args.align_thr,
        left_right_thr_px=args.lr_thr,
    )

    # Speed mode
    if args.low:
        linear_speed = params.low_linear_speed
        angular_speed = params.low_angular_speed
    else:
        linear_speed = params.linear_speed
        angular_speed = params.angular_speed

    # Mode and posture pre-setup (best-effort)
    set_autonomous_low_pose_if_possible()

    # Video capture (CPU, V4L2 backend)
    cap = cv.VideoCapture(args.camera, cv.CAP_V4L2)
    if not cap.isOpened():
        print(f"[ERROR] Cannot open camera {args.camera}")
        switch_back_to_manual_if_possible()
        sys.exit(2)

    # Center x for alignment
    frame_width = cap.get(cv.CAP_PROP_FRAME_WIDTH)
    if frame_width <= 0:
        frame_width = 640.0
    center_x = frame_width / 2.0

    mode, dictionary, detector_or_params = build_aruco_detector()

    print("[INFO] ArUco navigation started. Press Ctrl+C to exit.")
    try:
        while not rospy.is_shutdown():
            ok, frame = cap.read()
            if not ok:
                rate.sleep()
                continue

            if mode == "new":
                detector = detector_or_params
                corners, ids, _ = detector.detectMarkers(frame)
            else:
                params_legacy = detector_or_params
                corners, ids, _ = cv.aruco.detectMarkers(frame, dictionary, parameters=params_legacy)

            linear_cmd = 0.0
            angular_cmd = 0.0

            target_found = False
            if ids is not None and len(corners) > 0:
                for idx, marker_id in enumerate(np.squeeze(ids)):
                    if int(marker_id) == params.target_id:
                        target_found = True
                        target_corners = corners[idx]
                        cx = get_marker_center_x(target_corners)
                        side_px = compute_marker_side_length_px(target_corners)

                        # Alignment error (pixels)
                        err_px = cx - center_x

                        # If not well aligned, rotate in place (no forward motion)
                        if abs(err_px) > params.forward_center_thr_px:
                            # proportional control on angular speed
                            norm_err = float(err_px / center_x)  # ~[-1, 1]
                            angular_cmd = float(np.clip(norm_err, -1.0, 1.0)) * angular_speed
                            linear_cmd = 0.0
                        else:
                            # Aligned enough: go forward until close
                            if side_px >= params.stop_side_px:
                                linear_cmd = 0.0
                                angular_cmd = 0.0
                                print("[INFO] Arrived at target (stop threshold reached).")
                                publish_twist(cmd_pub, 0.0, 0.0)
                                return
                            else:
                                linear_cmd = linear_speed
                                angular_cmd = 0.0

                        if args.gui:
                            cv.aruco.drawDetectedMarkers(frame, [target_corners], np.array([[params.target_id]]))
                            cv.putText(frame, f"err_px={int(err_px)} side_px={int(side_px)}", (10, 30), cv.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
                        break

            # Publish velocity
            publish_twist(cmd_pub, linear_cmd, angular_cmd)

            # GUI display
            if args.gui:
                if ids is not None and len(corners) > 0 and not target_found:
                    # draw all for debugging
                    if mode == "new":
                        cv.aruco.drawDetectedMarkers(frame, corners, ids)
                    else:
                        cv.aruco.drawDetectedMarkers(frame, corners, ids)
                cv.line(frame, (int(center_x), 0), (int(center_x), frame.shape[0]), (255, 0, 0), 1)
                cv.imshow("camera", frame)
                key = cv.waitKey(1) & 0xFF
                if key == ord('q'):
                    break

            rate.sleep()
    except KeyboardInterrupt:
        pass
    finally:
        # Stop motion and cleanup
        try:
            publish_twist(cmd_pub, 0.0, 0.0)
        except Exception:
            pass
        cap.release()
        if args.gui:
            cv.destroyAllWindows()
        switch_back_to_manual_if_possible()
        print("[INFO] ArUco navigation finished.")


if __name__ == "__main__":
    main()


